'''
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Created on Jul 25, 2009

Version: 

@author: bryan
'''

import JT

import numpy as np
from math import atan2

class AbstractCalibration():
    """
    A base class for calibration objects.  This class should not be used directly,
    instead use the more specific InernalCalibration or ExternalCalibration classes.
    """

    def isValid(self):
        """
        Return True if calbration has been correctly set,
        return False if calibration has not been set or if contained errors.
        """
        return self._isValidDataBool

    def SetValid(self):
        self._isValidDataBool = True


class InternalCalibration(AbstractCalibration):
    """
    Internal Calibration parameters (focal point, offset, scale).

    If only one camera/fluoro is being used, only the internal calibration is
    necessary.  For multiple camera/fluoro setups, an external calibration also must be
    provided.
    """

    def __init__(self):
        """
        Read JointTrack calibration information and perform basic validation on
        the parameters.
        """
        self.cal_data = []

        self._oldFD = []
        self._oldXO = []
        self._oldYO = []
        self._oldPS = []

        self._FocalDistance = 1.0
        self._XOffset = 0.0
        self._YOffset = 0.0
        self._PixelScale = 1.0

        self._isValidDataBool = False


    def LoadFromFile(self,filename):
        """
        Load a JointTrack calibration file.
        """
        file = open(filename)
        raw_data = file.read()
        data = raw_data.split()

        if data[0] != 'JT_INTCALIB':
            print filename, " is not a valid JT internal calibration file, nothing done!"
            return

        #Save a copy of the old calibration in case it needs to be restored later
        self.oldFD = self._FocalDistance
        self.oldXO = self._XOffset
        self.oldYO = self._YOffset
        self.oldPS = self._PixelScale

        #Just save the relevant data to cal_data
        cal_data = data[1:5]

        #Check to make sure that the calibration appears good
        self.CheckCalibration(cal_data)
        print filename, cal_data

        if self._isValidDataBool == True:
            #save the new individual parameters
            self._FocalDistance = float(cal_data[0])
            self._XOffset = float(cal_data[1])
            self._YOffset = float(cal_data[2])
            self._PixelScale = float(cal_data[3])

        elif self._isValidDataBool == False:
            print "Calibration data was not numeric, nothing done"


    def GetCals(self):
        """
        Get the current calibration as a string.
        """

        cal_string = "Focal Distance: %.2f, X Offset: %.2f, Y Offset: %.2f, Pixel Scale: %.2f" \
            % (self._FocalDistance,self._XOffset,self._YOffset,self._PixelScale)
        return cal_string

    def GetCald(self):
        """
        Get the current calibration as a dictionary.
        """

        cal_dict = {'Focal Distance': self._FocalDistance,
        'X Offset': self._XOffset,'Y Offset':self._YOffset,'Pixel Scale':self._PixelScale}

        return cal_dict

    def PrintCals(self):
        """
        Print out a pretty looking version of the current calibration.
        """
        print "Current Calibration"
        print "   Focal Distance: ", self._FocalDistance
        print "   X Offset: ", self._XOffset
        print "   Y Offset: ", self._YOffset
        print "   Pixel Scale: ", self._PixelScale

        if self._isValidDataBool:
            print "Calibration is valid: "
        else:
            print "Calibration is NOT valid: "

        if self._calErrors != []:
            print "Calibration Errors: "
            print     self._calErrors
        else:
            print "No Calibration Errors."


    def SetFocalDistance(self,dist):
        self._FocalDistance = float(dist)

    def SetXOffset(self,offset):
        self._XOffset = offset

    def SetYOffset(self,offset):
        self._YOffset = offset

    def SetPixelScale(self,scale):
        self._PixelScale = scale


    def CheckCalibration(self, cal_data):
        self._calErrors = []
        self._isValidDataBool = False

        try:
            assert(abs(float(cal_data[0]))>=0)
        except:
            self._calErrors.append("Test Failure! Focal Distance is NOT valid.")

        try:
            assert(abs(float(cal_data[1]))>=0)
        except:
            self._calErrors.append("Test Failure! X Offset is NOT valid.")

        try:
            assert(abs(float(cal_data[2]))>=0)
        except:
            self._calErrors.append("Test Failure! Y Offset is NOT valid.")

        try:
            assert(abs(float(cal_data[3]))>=0)
        except:
            self._calErrors.append("Test Failure! Pixel Scale is NOT valid.")


        if self._calErrors == []:
            self._isValidDataBool = True
        else:
            print self._calErrors


class ExternalCalibration(AbstractCalibration):
    """
    External Calibration.  Work in progress.
    """

    def __init__(self):
        """
        """
        self._pSource = []
        self._pDetector = []
        self._pVolume = []
        self._centerDetector = []
        self._isValidDataBool = False
        self.volumeTransform = None

    def LoadConsolidatedCalibration(self,filename):
        """DEPRECATED.  Use LoadCalFile.
        """
        print "LoadConsolidatedCalibration is DEPRECATED.  Use LoadCalFile."
        self.LoadCalFile(filename)
    
    def LoadCalFile(self,filename):
        """
        Camera postion, volume position, detector position.  Consolidated calibration,
        one file consisting of all the calibration data from *.txt, *.rad, *.sta files.

        Other methods are provided for reading each of these files separately
        """
        try:
            file = open(filename)
        except IOError:
            raise

        raw_data = file.read()
        data = raw_data.split()

        if data[0] != 'JT_EXTCALIB':
            print filename, " is not a valid JT exernal calibration file, nothing done!"
            return

        #Just save the relevant data to cal_data
        # Position of the source in world coordinates
        pSource = np.array([float(data[2]),float(data[3]),float(data[4])])
        # Position of the Detector in world coordinates
        pDetector = np.array([float(data[6]),float(data[7]),float(data[8])])
        # Position of the volume in world coordinates, this is not relevant for
        # calibration and would not be known in real life situations
        pVolume = np.array([float(data[10]),float(data[11]),float(data[12])])
        # Size of the image (in pixels) projected on the detector
        fixedImageSize = np.array([float(data[14]),float(data[15])])
        # Pixel size of detector (in world units)
        fixedImagePixelScale = np.array([float(data[17]),float(data[18])])
        # Size of the Volume(in pixels), not relevant for calibration
        ImageSize3D = np.array([float(data[20]),float(data[21]),float(data[22])])
        # Pixel size of volume (in world units)
        volumePixelScale = np.array([float(data[24]),float(data[25]),float(data[26])])

        # Transform of the Detector relative to world.
        DetectorT = np.array([[float(data[28]),float(data[29]),float(data[30]),float(data[31])],
                            [float(data[32]),float(data[33]),float(data[34]),float(data[35])],
                            [float(data[36]),float(data[37]),float(data[38]),float(data[39])],
                            [float(data[40]),float(data[41]),float(data[42]),float(data[43])]])
        

        # The position of the detector is relative to its origin, which is
        # located at the top left corner.  To display the image properly,
        # the camera mmust be pointed at the center of the detector.  Therefore,
        # the focal point must be shifted by 1/2 the image size (in world
        # units).  The direction of the shift is determined by the orientation
        # of the detector in the world coordinate system.

        # offsetDetector is the amount of the shift (homogenous vector of length 4)
        # along the detector x and y axes.
        offsetDetector = np.ones(4)
        offsetDetector[2] = 0
        offsetDetector[0:2] = 0.5 * fixedImageSize * fixedImagePixelScale

        # The center of the detector is found by using matrix multiplication
        # between the detector transform (wrt World) and the offset (wrt detector)
        centerDetector = np.dot(DetectorT,offsetDetector)

        # Check to make sure that the calibration appears good
        # self.CheckCalibration(cal_data)
        # Short Circuit the calibration checking method
        self._isValidDataBool = True

        if self._isValidDataBool == True:
            #save the new individual parameters
            self._pSource = pSource
            self._pDetector = pDetector
            self._pVolume = pVolume
            self.fixedImageSize = fixedImageSize
            self.fixedImagePixelScale = fixedImagePixelScale
            self.volumeSize = ImageSize3D
            self.volumePixelScale = volumePixelScale
            #self._FocalDistance = mag(pDetector-pSource)
            self._FocalDistance = JT.utils.mag(centerDetector[0:3]-pSource)
            self._DetectorT = DetectorT
            self._centerDetector = centerDetector[0:3]

        elif self._isValidDataBool == False:
            print "Calibration data was not numeric, nothing done"
            #Don't have to restore old data, because they should not have been changed


        self.ComputeViewAngle()


    def LoadStaFile(self,filename):
        """
        Load *.sta cal file.  Which is specifies the position, image parameters,
        and orientation of the volume.
        """
        try:
            file = open(filename)
        except IOError:
            raise

        raw_data = file.read()
        data = raw_data.split()

        self.volumeSize = np.array([float(data[1]),float(data[2]),float(data[3])])
        # Pixel size of volume (in world units)
        self.volumePixelScale = np.array([float(data[5]),float(data[6]),float(data[7])])

        # Transform of the Detector relative to world.
        self.volumeTransform = np.array([[float(data[9]),float(data[10]),float(data[11]),float(data[12])],
                            [float(data[13]),float(data[14]),float(data[15]),float(data[16])],
                            [float(data[17]),float(data[18]),float(data[19]),float(data[20])],
                            [float(data[21]),float(data[22]),float(data[23]),float(data[24])]])

#        JT.logging.debug("Sta file loaded.")
#        JT.logging.debug(" 3D Image Size: %s" % self.volumeSize)
#        JT.logging.debug(" 3D Image Pixel Scale: %s" % self.volumePixelScale)
#        JT.logging.debug(" 3D Image Transform: %s" % self.volumeTransform)


    def LoadRadFile(self,filename):
        """
        Load *.rad cal file.  Which is specifies the position, image parameters,
        and orientation of the fluoro image.
        """
        try:
            file = open(filename)
        except IOError:
            raise

        raw_data = file.read()
        data = raw_data.split()

        self.fixedImageSize = np.array([float(data[1]),float(data[2])])
        # Pixel size of volume (in world units)
        self.fixedImagePixelScale = np.array([float(data[4]),float(data[5])])

        # Transform of the Detector relative to world.
        self._DetectorT = np.array([[float(data[7]),float(data[8]),float(data[9]),float(data[10])],
                            [float(data[11]),float(data[12]),float(data[13]),float(data[14])],
                            [float(data[15]),float(data[16]),float(data[17]),float(data[18])],
                            [float(data[19]),float(data[20]),float(data[21]),float(data[22])]])

        # The position of the detector is relative to its origin, which is
        # located at the top left corner.  To display the image properly,
        # the camera mmust be pointed at the center of the detector.  Therefore,
        # the focal point must be shifted by 1/2 the image size (in world
        # units).  The direction of the shift is determined by the orientation
        # of the detector in the world coordinate system.

        # offsetDetector is the amount of the shift (homogenous vector of length 4)
        # along the detector x and y axes.
        offsetDetector = np.ones(4)
        offsetDetector[2] = 0
        offsetDetector[0:2] = 0.5 * self.fixedImageSize * self.fixedImagePixelScale

        # The center of the detector is found by using matrix multiplication
        # between the detector transform (wrt World) and the offset (wrt detector)
        self._centerDetector = np.dot(self._DetectorT,offsetDetector)

        #if the position of the fluoro source has been set, calculate the focal distance
        if self._pSource != []:
            self._FocalDistance = JT.utils.mag(self._centerDetector[0:3]-self._pSource)

        JT.logging.debug("Rad file loaded.")
        JT.logging.debug(" 2D Image Size: %s" % self.fixedImageSize)
        JT.logging.debug(" 2D Image Pixel Scale: %s" % self.fixedImagePixelScale)
        JT.logging.debug(" 2D Image Transform: %s" % self._DetectorT)
        JT.logging.debug(" Center Detector (world coordinates): %s" % self._centerDetector)

    def LoadPSourceFile(self,filename):
        """
        Load a pSource*.txt file which specifies the position of the fluoro source
        (camera) in world coordinates.
        """
        try:
            file = open(filename)
        except IOError:
            raise

        raw_data = file.read()
        data = raw_data.split()

        self._pSource = np.array([float(data[0]),float(data[1]),float(data[2])])

        #if the detector has been calibrated, calculate the focal distance
        if self._centerDetector != []:
            self._FocalDistance = JT.utils.mag(self._centerDetector[0:3]-self._pSource)

        JT.logging.debug("Source Position (pSource): %s" % self._pSource)

    def CheckCalibration(self, cal_data):
        """
        TODO: Not implemented

        Need to update this function.  Check that each element is a number
        and that there are 3 elements in each position variable.
        """
        self._calErrors = []
        self._isValidDataBool = False

        try:
            assert(abs(float(cal_data[0]))>=0)
        except:
            self._calErrors.append("Test Failure! Focal Distance is NOT valid.")

        try:
            assert(abs(float(cal_data[1]))>=0)
        except:
            self._calErrors.append("Test Failure! X Offset is NOT valid.")

        try:
            assert(abs(float(cal_data[2]))>=0)
        except:
            self._calErrors.append("Test Failure! Y Offset is NOT valid.")

        try:
            assert(abs(float(cal_data[3]))>=0)
        except:
            self._calErrors.append("Test Failure! Pixel Scale is NOT valid.")


        if self._calErrors == []:
            self._isValidDataBool = True
        else:
            print self._calErrors


    def ComputeViewAngle(self):
        """
        This function will calculate the view angle required to fully visualize
        the the desired image.  ImageSize should be the size of the displayed
        image in world dimensions (eg. mm).

        The calibration must already be set before this method is used!
        """
        if self.fixedImagePixelScale == []:
            "ERROR: 2D Pixel Scale, 2D Image Size, and Focal Distance must be set to calculate view angle"
            return
        if self.fixedImageSize == []:
            "ERROR: 2D Pixel Scale, 2D Image Size, and Focal Distance must be set to calculate view angle"
            return
        if self._FocalDistance == []:
            "ERROR: 2D Pixel Scale, 2D Image Size, and Focal Distance must be set to calculate view angle"
            return

        h = self.fixedImagePixelScale[1] * self.fixedImageSize[1] #height of render window in mm
        r2d = 180.0/np.pi

        view_angle = 2.0 * r2d * atan2((h / 2.0), self._FocalDistance)
        self._viewAngle = view_angle


    def ConvertNumpyToVTKArray(self,npArray):
        """
        Numpy array is returned from calibration, but a vtk matrix is required
        for setting pose of a vtk volume.
        """
        vtkMatrix = JT.vtk.vtkMatrix4x4()

        if (npArray.size != 16) or (len(npArray) != 4):
            print "Error: numpy array is not 4x4, cannot convert to VTK matrix"
            return

        for row in range(len(npArray)):
            for col in range(len(npArray[row])):
                vtkMatrix.SetElement(row,col,npArray[row][col])

        return vtkMatrix

    def PrintCal(self):
        """
        """

        print filename
        print " Source: ", self._pSource
        print " Detector: ", self._pDetector
        print " Detector Transform: ", self._DetectorT
        print " Center Detector (in World): ", self._centerDetector
        print " Volume: ", self._pVolume
        print " 2D Size: ", self.fixedImageSize
        print " 2D Pixel Scale: ", self.fixedImagePixelScale
        print " 3D Size: ", self.volumeSize
        print " 3D Pixel Scale: ", self.volumePixelScale
        print " View Angle: ", self._viewAngle
        print " FocalDistance: ", self._FocalDistance
